package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"strconv"
	"syscall"
	"time"

	"network-firewall-go/internal/cert"
	"network-firewall-go/internal/database"
	"network-firewall-go/internal/proxy"
	"network-firewall-go/internal/rules"
	"network-firewall-go/internal/systemproxy"
	"network-firewall-go/internal/types"
	"network-firewall-go/internal/utils"
	"network-firewall-go/internal/websocket"
)

var (
	port       = flag.Int("port", 12300, "HTTP代理端口")
	apiPort    = flag.Int("api-port", 12280, "API服务端口")
	wsPort     = flag.Int("ws-port", 12290, "WebSocket端口")
	intercept  = flag.Bool("intercept-https", true, "是否拦截HTTPS请求")
	certDir    = flag.String("cert-dir", "certs", "证书目录")
	configFile = flag.String("config", "configs/config.json", "配置文件路径")
	dbPath     = flag.String("db", "network-firewall.db", "数据库文件路径")
	debug      = flag.Bool("debug", false, "启用调试模式")
	logLevel   = flag.String("log-level", "info", "日志级别")
	autoProxy  = flag.Bool("auto-proxy", true, "自动设置系统代理")
)

// App 应用程序
type App struct {
	config      *types.ProxyConfig
	certManager *cert.Manager
	database    *database.Manager
	ruleEngine  *rules.Engine
	wsServer    *websocket.Server
	proxyServer *proxy.Server
	systemProxy *systemproxy.Manager
	apiServer   *http.Server
}

func main() {
	flag.Parse()

	fmt.Println("🚀 Network Firewall Go - 高性能网络防火墙")
	fmt.Println("========================================")

	// 创建应用程序
	app, err := NewApp()
	if err != nil {
		log.Fatalf("❌ 创建应用程序失败: %v", err)
	}

	// 设置紧急清理函数（确保程序异常退出时也能清理代理）
	defer func() {
		if r := recover(); r != nil {
			fmt.Printf("🚨 程序异常退出，执行紧急清理: %v\n", r)
			if *autoProxy && app.systemProxy != nil {
				app.systemProxy.ForceDisableProxy()
			}
		}
	}()

	// 启动应用程序
	if err := app.Start(); err != nil {
		// 启动失败时也要清理可能已设置的代理
		if *autoProxy && app.systemProxy != nil {
			fmt.Println("🚨 启动失败，清理可能的代理设置...")
			app.systemProxy.ForceDisableProxy()
		}
		log.Fatalf("❌ 启动应用程序失败: %v", err)
	}

	// 等待信号
	app.WaitForSignal()

	// 停止应用程序
	if err := app.Stop(); err != nil {
		log.Printf("❌ 停止应用程序失败: %v", err)
		// 如果正常停止失败，强制清理代理
		if *autoProxy && app.systemProxy != nil {
			fmt.Println("🚨 正常停止失败，执行强制清理...")
			app.systemProxy.ForceDisableProxy()
		}
	}

	fmt.Println("👋 再见！")
}

// NewApp 创建新的应用程序
func NewApp() (*App, error) {
	// 加载配置
	config := &types.ProxyConfig{
		Port:           *port,
		InterceptHTTPS: *intercept,
		CertPath:       filepath.Join(*certDir, "root-ca.crt"),
		KeyPath:        filepath.Join(*certDir, "root-ca.key"),
		DatabasePath:   *dbPath,
		LogLevel:       *logLevel,
	}

	// 如果配置文件存在，则加载配置文件
	if _, err := os.Stat(*configFile); err == nil {
		if err := loadConfigFromFile(*configFile, config); err != nil {
			return nil, fmt.Errorf("加载配置文件失败: %w", err)
		}
	}

	// 创建证书管理器
	certManager, err := cert.NewManager(*certDir)
	if err != nil {
		return nil, fmt.Errorf("创建证书管理器失败: %w", err)
	}

	// 创建数据库管理器
	db, err := database.NewManager(config.DatabasePath)
	if err != nil {
		return nil, fmt.Errorf("创建数据库管理器失败: %w", err)
	}

	// 创建规则引擎
	ruleEngine := rules.NewEngine(db)

	// 创建WebSocket服务器
	wsServer := websocket.NewServer()

	// 创建代理服务器
	proxyServer := proxy.NewServer(config.Port, db, ruleEngine, config, certManager, wsServer)

	// 创建系统代理管理器
	systemProxy := systemproxy.NewManager()

	return &App{
		config:      config,
		certManager: certManager,
		database:    db,
		ruleEngine:  ruleEngine,
		wsServer:    wsServer,
		proxyServer: proxyServer,
		systemProxy: systemProxy,
	}, nil
}

// Start 启动应用程序
func (app *App) Start() error {
	fmt.Println("🔧 正在启动各个组件...")

	// 启动WebSocket服务器
	app.wsServer.Start()

	// 加载规则
	if err := app.ruleEngine.LoadRules(); err != nil {
		fmt.Printf("⚠️ 加载规则失败: %v\n", err)
	}

	// 启动代理服务器
	if err := app.proxyServer.Start(); err != nil {
		return fmt.Errorf("启动代理服务器失败: %w", err)
	}

	// 启动API服务器
	if err := app.startAPIServer(); err != nil {
		return fmt.Errorf("启动API服务器失败: %w", err)
	}

	// 安装根证书到系统信任存储
	if err := app.installRootCertificate(); err != nil {
		fmt.Printf("⚠️ 安装根证书失败: %v\n", err)
	}

	// 启用系统代理（如果启用）
	if *autoProxy {
		fmt.Println("📡 启用系统代理...")
		if err := app.systemProxy.Enable(app.config.Port); err != nil {
			fmt.Printf("⚠️ 启用系统代理失败: %v\n", err)
		}
	}

	fmt.Printf("✅ 所有组件启动完成\n")
	fmt.Printf("🌐 API服务: http://0.0.0.0:%d\n", *apiPort)
	fmt.Printf("📡 WebSocket: ws://0.0.0.0:%d/ws\n", *apiPort)

	return nil
}

// Stop 停止应用程序
func (app *App) Stop() error {
	fmt.Println("🛑 正在停止各个组件...")

	// 优先恢复系统代理，确保网络连接
	if *autoProxy {
		fmt.Println("🔧 正在恢复系统代理设置...")
		if err := app.systemProxy.Disable(); err != nil {
			fmt.Printf("⚠️ 恢复系统代理失败: %v\n", err)
			// 即使失败，也尝试强制清理
			fmt.Println("🚨 尝试强制清理系统代理...")
			if forceErr := app.systemProxy.ForceDisableProxy(); forceErr != nil {
				fmt.Printf("❌ 强制清理也失败: %v\n", forceErr)
			}
		}
	}

	// 停止API服务器
	if app.apiServer != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		app.apiServer.Shutdown(ctx)
	}

	// 停止代理服务器
	if err := app.proxyServer.Stop(); err != nil {
		fmt.Printf("⚠️ 停止代理服务器失败: %v\n", err)
	}

	// 关闭数据库
	if err := app.database.Close(); err != nil {
		fmt.Printf("⚠️ 关闭数据库失败: %v\n", err)
	}

	fmt.Println("✅ 所有组件已停止")
	return nil
}

// startAPIServer 启动API服务器
func (app *App) startAPIServer() error {
	mux := http.NewServeMux()

	// WebSocket端点
	mux.HandleFunc("/ws", app.wsServer.HandleWebSocket)

	// API端点
	mux.HandleFunc("/api/status", app.handleGetStatus)
	mux.HandleFunc("/api/proxy/start", app.handleStartProxy)
	mux.HandleFunc("/api/proxy/stop", app.handleStopProxy)
	mux.HandleFunc("/api/transactions", app.handleGetTransactions)
	mux.HandleFunc("/api/transactions/", app.handleGetTransaction)
	mux.HandleFunc("/api/rules", app.handleRules)
	mux.HandleFunc("/api/rules/", app.handleRule)
	mux.HandleFunc("/api/system/proxy", app.handleSystemProxy)
	mux.HandleFunc("/api/cert/root", app.handleRootCert)

	// 静态文件服务（如果存在前端文件）
	if _, err := os.Stat("web"); err == nil {
		mux.Handle("/", http.FileServer(http.Dir("web")))
	}

	app.apiServer = &http.Server{
		Addr:    fmt.Sprintf(":%d", *apiPort),
		Handler: enableCORS(mux),
	}

	go func() {
		if err := app.apiServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Printf("❌ API服务器错误: %v", err)
		}
	}()

	return nil
}

// WaitForSignal 等待信号
func (app *App) WaitForSignal() {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	fmt.Println("✅ 服务器运行中，按 Ctrl+C 停止...")
	<-sigChan
	fmt.Println("\n📡 收到停止信号...")
}

// API处理函数

// handleGetStatus 获取状态
func (app *App) handleGetStatus(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	status := app.proxyServer.GetStatus()
	status.Connections = app.wsServer.GetClientCount()

	respondJSON(w, status)
}

// handleStartProxy 启动代理
func (app *App) handleStartProxy(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	if err := app.proxyServer.Start(); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	status := app.proxyServer.GetStatus()
	app.wsServer.BroadcastStatus(status)

	respondJSON(w, map[string]string{"message": "代理服务器已启动"})
}

// handleStopProxy 停止代理
func (app *App) handleStopProxy(w http.ResponseWriter, r *http.Request) {
	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	if err := app.proxyServer.Stop(); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	status := app.proxyServer.GetStatus()
	app.wsServer.BroadcastStatus(status)

	respondJSON(w, map[string]string{"message": "代理服务器已停止"})
}

// handleGetTransactions 获取事务列表
func (app *App) handleGetTransactions(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	// 解析查询参数
	limit := 100
	offset := 0

	if limitStr := r.URL.Query().Get("limit"); limitStr != "" {
		if l, err := strconv.Atoi(limitStr); err == nil && l > 0 {
			limit = l
		}
	}

	if offsetStr := r.URL.Query().Get("offset"); offsetStr != "" {
		if o, err := strconv.Atoi(offsetStr); err == nil && o >= 0 {
			offset = o
		}
	}

	transactions, err := app.database.GetTransactions(limit, offset)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// 获取总数
	stats, err := app.database.GetStats()
	var total int64 = 0
	if err == nil {
		if totalTransactions, ok := stats["total_transactions"].(int64); ok {
			total = totalTransactions
		}
	}

	// 返回与前端期望一致的格式
	response := map[string]interface{}{
		"data":  transactions,
		"total": total,
	}

	respondJSON(w, response)
}

// handleGetTransaction 获取单个事务
func (app *App) handleGetTransaction(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	id := r.URL.Path[len("/api/transactions/"):]
	if id == "" {
		http.Error(w, "Transaction ID required", http.StatusBadRequest)
		return
	}

	transaction, err := app.database.GetTransactionByID(id)
	if err != nil {
		http.Error(w, err.Error(), http.StatusNotFound)
		return
	}

	respondJSON(w, transaction)
}

// handleRules 处理规则相关请求
func (app *App) handleRules(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		rules := app.ruleEngine.GetRules()
		respondJSON(w, rules)
	case "POST":
		var rule types.Rule
		if err := json.NewDecoder(r.Body).Decode(&rule); err != nil {
			http.Error(w, "Invalid JSON", http.StatusBadRequest)
			return
		}

		rule.ID = utils.GenerateID()
		if err := app.ruleEngine.AddRule(&rule); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		app.wsServer.BroadcastRuleUpdate(&rule)
		respondJSON(w, rule)
	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

// handleRule 处理单个规则请求
func (app *App) handleRule(w http.ResponseWriter, r *http.Request) {
	id := r.URL.Path[len("/api/rules/"):]
	if id == "" {
		http.Error(w, "Rule ID required", http.StatusBadRequest)
		return
	}

	switch r.Method {
	case "DELETE":
		if err := app.ruleEngine.DeleteRule(id); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		respondJSON(w, map[string]string{"message": "规则已删除"})
	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

// handleSystemProxy 处理系统代理请求
func (app *App) handleSystemProxy(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		config, err := app.systemProxy.GetCurrentConfig()
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		respondJSON(w, config)
	case "POST":
		if err := app.systemProxy.Enable(app.config.Port); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		respondJSON(w, map[string]string{"message": "系统代理已启用"})
	case "DELETE":
		if err := app.systemProxy.Disable(); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		respondJSON(w, map[string]string{"message": "系统代理已禁用"})
	default:
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
	}
}

// handleRootCert 处理根证书请求
func (app *App) handleRootCert(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	certPEM, err := app.certManager.GetRootCAPEM()
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/x-pem-file")
	w.Header().Set("Content-Disposition", "attachment; filename=root-ca.crt")
	w.Write(certPEM)
}

// 工具函数

// respondJSON 响应JSON
func respondJSON(w http.ResponseWriter, data interface{}) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(data)
}

// enableCORS 启用CORS
func enableCORS(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")

		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}

		next.ServeHTTP(w, r)
	})
}

// loadConfigFromFile 从文件加载配置
func loadConfigFromFile(filename string, config *types.ProxyConfig) error {
	data, err := os.ReadFile(filename)
	if err != nil {
		return err
	}

	return json.Unmarshal(data, config)
}

// installRootCertificate 安装根证书到系统信任存储
func (app *App) installRootCertificate() error {
	// 获取根证书路径
	rootCertPath := filepath.Join(*certDir, "root-ca.crt")

	// 检查证书文件是否存在
	if _, err := os.Stat(rootCertPath); os.IsNotExist(err) {
		return fmt.Errorf("根证书文件不存在: %s", rootCertPath)
	}

	// 安装证书到系统信任存储
	if err := app.systemProxy.InstallCertificate(rootCertPath); err != nil {
		return fmt.Errorf("安装根证书失败: %w", err)
	}

	return nil
}
